# -*- coding: utf-8 -*-#

# -------------------------------------------------------------------------------
# Name:         zx_inference
# Description:
# Author:       zx
# Date:         2020/10/13
# -------------------------------------------------------------------------------
import argparse
import fnmatch
import os

import math
import time

import cv2
import os.path as osp
import glob
import numpy as np
from shapely.geometry import Polygon
import pyclipper

from DetectionIoUEvaluator import DetectionIoUEvaluator
# from model import dbnet
from zx_model import dbnet


def resize_image(image, label_info, image_short_side=736):
    height, width, _ = image.shape
    print(height, width, _)
    if height < width:
        new_height = image_short_side
        new_width = int(math.ceil(new_height / height * width / 32) * 32)
    else:
        new_width = image_short_side
        new_height = int(math.ceil(new_width / width * height / 32) * 32)
    resized_img = cv2.resize(image, (new_width, new_height))
    print(new_width, new_height)

    for i in range(len(label_info)):
        label_info[i]["points"] = [tuple([e[0] * new_width / width, e[1] * new_height / height]) for e in
                                   label_info[i]["points"]]
        # label_info[i]["points"] = [tuple([e[0] * new_height / height, e[1] * new_width / width]) for e in
        #                            label_info[i]["points"]]
    return resized_img, label_info


def box_score_fast(bitmap, _box):
    h, w = bitmap.shape[:2]
    box = _box.copy()
    xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
    xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int), 0, w - 1)
    ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int), 0, h - 1)
    ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int), 0, h - 1)

    mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
    box[:, 0] = box[:, 0] - xmin
    box[:, 1] = box[:, 1] - ymin
    cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
    return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]


def unclip(box, unclip_ratio=1.5):
    poly = Polygon(box)
    distance = poly.area * unclip_ratio / poly.length
    offset = pyclipper.PyclipperOffset()
    offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
    expanded = np.array(offset.Execute(distance))
    return expanded


def get_mini_boxes(contour):
    bounding_box = cv2.minAreaRect(contour)
    points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])

    index_1, index_2, index_3, index_4 = 0, 1, 2, 3
    if points[1][1] > points[0][1]:
        index_1 = 0
        index_4 = 1
    else:
        index_1 = 1
        index_4 = 0
    if points[3][1] > points[2][1]:
        index_2 = 2
        index_3 = 3
    else:
        index_2 = 3
        index_3 = 2

    box = [points[index_1], points[index_2],
           points[index_3], points[index_4]]
    return box, min(bounding_box[1])


def polygons_from_bitmap(pred, bitmap, dest_width, dest_height, max_candidates=100, box_thresh=0.7):
    pred = pred[..., 0]
    bitmap = bitmap[..., 0]
    height, width = bitmap.shape
    boxes = []
    scores = []

    _, contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    for contour in contours[:max_candidates]:
        epsilon = 0.01 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        points = approx.reshape((-1, 2))
        if points.shape[0] < 4:
            continue
        score = box_score_fast(pred, points.reshape(-1, 2))
        if box_thresh > score:
            continue
        if points.shape[0] > 2:
            box = unclip(points, unclip_ratio=2.0)
            if len(box) > 1:
                continue
        else:
            continue
        box = box.reshape(-1, 2)
        _, sside = get_mini_boxes(box.reshape((-1, 1, 2)))
        if sside < 5:
            continue

        box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
        box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
        boxes.append({
            "points": [tuple(e) for e in box.tolist()],
            'text': 1234,
            'ignore': False,
        })
        scores.append(score)
    return boxes, scores


def load_each_image_lable(lable_path):
    lines = []
    reader = open(lable_path, 'r').readlines()
    for line in reader:
        item = {}
        parts = line.strip().split(',')
        label = parts[-1]
        num_points = math.floor((len(parts) - 1) / 2) * 2
        poly = np.array(list(map(float, parts[:num_points]))).reshape((-1, 2)).tolist()
        item['points'] = [tuple(e) for e in poly]
        item['text'] = label
        item['ignore'] = False
        lines.append(item)
    return lines


def parser_args():
    parser = argparse.ArgumentParser(description='choose checkpoints and epoch')
    parser.add_argument('--checkpoints', type=int, help='Resume from checkpoint', default=1)
    parser.add_argument('--day', type=int, help='Resume from day', default=22)
    parser.add_argument('--epoch', type=int, help='Resume from epoch', default=48)
    args = vars(parser.parse_args())

    # for file in os.listdir('/data1/yhm/checkpoints/2020-09-27/'):
    #     if fnmatch.fnmatch(file, 'db_{}_*'.format(args["epoch"])):
    #         return True, '/data1/yhm/checkpoints/2020-09-27/' + file
    # return False, None
    for file in os.listdir('./models/'):
        if fnmatch.fnmatch(file, 'db_{}_*'.format(args["epoch"])):
            return True, './models/' + file
    return False, None


if __name__ == '__main__':
    find, file = parser_args()
    print(find, file)
    # file = "/data1/yhm/checkpoints/2020-09-27/db_127_4.9566_4.9575.h5"
    if not find:
        print("not find")
        exit(0)
    evaluator = DetectionIoUEvaluator()
    gts = []
    preds = []

    mean = np.array([104.00698793, 116.66876762, 122.67891434])

    i = 0
    _, model = dbnet()

    model.load_weights(file, by_name=True, skip_mismatch=True)
    for image_path in glob.glob(osp.join('datasets/total_text/test_images', '*.jpg'))[0:4]:
        print(image_path)
        image = cv2.imread(image_path)
        src_image = image.copy()
        h, w = image.shape[:2]
        label_info = load_each_image_lable("datasets/total_text/test_gts/{}.txt".format(image_path[32:]))
        image, label_info = resize_image(image, label_info)
        # cv2.imshow('resize_image', image)
        image = image.astype(np.float32)
        image -= mean
        # image = image / 255.0
        # cv2.imshow('mean', image.astype(np.float32))
        image_input = np.expand_dims(image, axis=0)
        start = time.time()
        p = model.predict(image_input)[0]
        time_cost = (time.time() - start)
        print("Inference speed: %fms, FPS: %f" % (time_cost * 1000, 1 / time_cost))
        bitmap = p > 0.3
        boxes, scores = polygons_from_bitmap(p, bitmap, w, h, box_thresh=0.5)
        # print(boxes)
        # print(scores)
        gts.append(label_info)
        preds.append(boxes)
        # print(boxes)
        # cv2.imshow('p', (p * 255).astype(np.uint8))
        # cv2.imshow('bitmap', (bitmap * 255).astype(np.uint8))
        # src_image2 = src_image.copy()
        # for box in boxes:
        #     cv2.drawContours(src_image, [np.array(box["points"])], -1, (0, 255, 0), 2)
        # cv2.namedWindow('image', cv2.WINDOW_NORMAL)
        # cv2.imshow('image', src_image)
        #
        # for box in label_info:
        #     cv2.drawContours(src_image2, [np.array(box["points"]).astype(np.int)], -1, (0, 255, 0), 2)
        # cv2.namedWindow('image2', cv2.WINDOW_NORMAL)
        # cv2.imshow('image2', src_image2)
        # cv2.waitKey(0)
        # image_fname = osp.split(image_path)[-1]
        # cv2.imwrite('test/' + image_fname, src_image)

    results = []
    for gt, pred in zip(gts, preds):
        print(gt)
        print(pred)
        print(">>>>>>>>..")
        results.append(evaluator.evaluate_image(gt, pred))
    metrics = evaluator.combine_results(results)
    fmeasure_score = 2 * metrics["precision"] * metrics["recall"] / \
                     (metrics["precision"] + metrics["recall"] + 1e-8)
    metrics["fmeasure"] = fmeasure_score
    print(metrics)
